import os
import sys
sys.path.append(os.path.abspath(os.path.dirname(__file__) + '/' + '../..'))

import pandas as pd
import numpy as np

from python.tools import (
    clean_folder
)

################
## Parameters ##
################

input_folder = './get_graphs/input'
output_folder = './get_graphs/output/MSE_decomposition'
max_L = 12 # Number of bias coefficients used

########################
## Consensus forecast ##
########################

clean_folder(output_folder)
res = []

# Read in composite bias coefficients 
# estimated using consensus-level dataset
df_bias_coefs = pd.read_csv('{}/IRF_SPF_consensus.txt'.format(input_folder), 
    sep = '\t')
mask = df_bias_coefs['lag'] <= max_L
df_bias_coefs = df_bias_coefs.loc[mask, ]

# Read in consensus dataset
df_cons = pd.read_csv('{}/consensus_dataset.csv'.format(input_folder))
df_cons['Efor_Step2'] = df_cons['Realiz1'] - df_cons['SPFfor_Step2']

# Perform the MSE decomposition
MSE_cons = (df_cons['Efor_Step2'] ** 2).mean()
uncond_bias_cons = df_cons['Efor_Step2'].mean()
var_eps = (MSE_cons - uncond_bias_cons ** 2) / (1 + (df_bias_coefs['IRF'] ** 2).sum())
frac_irreducible_cons = var_eps / MSE_cons
frac_uncond_bias_cons = uncond_bias_cons ** 2 / MSE_cons
frac_aggregate_cons   = 1 - frac_irreducible_cons - frac_uncond_bias_cons

# Save results
res.append({'level': 'consensus',
            'frac_uncond_bias': frac_uncond_bias_cons,
            'frac_aggregate': frac_aggregate_cons,
            'frac_idiosyncratic': 0.0,
            'frac_irreducible': frac_irreducible_cons,
            'MSE': MSE_cons})

##########################
## Individual forecasts ##
##########################

# Read in individual dataset
df_ind = pd.read_csv('{}/individual_dataset.csv'.format(input_folder))
df_ind['Efor_Step2'] = df_ind['Realiz1'] - df_ind['SPFfor_Step2']

# Calculate weights
df_weight = df_ind.groupby('ID').count()['Efor_Step2'].reset_index()
df_weight['weight'] = df_weight['Efor_Step2'] / df_weight['Efor_Step2'].sum()
df_weight = df_weight[['ID', 'weight']]

# Calculated weighted-average MSE
df_ind['Efor_Step2_sq'] = df_ind['Efor_Step2'] ** 2
df_temp = df_ind.groupby('ID').mean()['Efor_Step2_sq'].reset_index()
df_temp = pd.merge(df_temp, df_weight, on = 'ID', how = 'left')
MSE_ind = (df_temp['Efor_Step2_sq'] * df_temp['weight']).sum()

# Calculate weighted-average unconditional bias
df_temp = df_ind.groupby('ID').mean()['Efor_Step2'].reset_index()
df_temp['bias_sq'] = df_temp['Efor_Step2'] ** 2
df_temp = pd.merge(df_temp, df_weight, on = 'ID', how = 'left')
uncond_bias_ind = (df_temp['bias_sq'] * df_temp['weight']).sum()

# Perform the MSE decomposition
frac_irreducible_ind   = var_eps / MSE_ind
frac_uncond_bias_ind   = uncond_bias_ind ** 2 / MSE_ind
frac_aggregate_ind     = (frac_aggregate_cons * MSE_cons) / MSE_ind
frac_idiosyncratic_ind = 1 - frac_irreducible_ind - frac_uncond_bias_ind - frac_aggregate_ind

# Save results
res.append({'level': 'individual',
            'frac_uncond_bias': frac_uncond_bias_ind,
            'frac_aggregate': frac_aggregate_ind,
            'frac_idiosyncratic': frac_idiosyncratic_ind,
            'frac_irreducible': frac_irreducible_ind,
            'MSE': MSE_ind})

######################
## Get LaTeX output ##
######################

res = pd.DataFrame(res)

# Convert to percent and pretify
for col in res.columns:
    if col != 'level':
        if col != 'MSE':
            res[col] = res[col].apply(lambda x: '{:.2f}%'.format(100 * x))
        elif col == 'MSE':
            res[col] = res[col].apply(lambda x: '{:.2f}'.format(x))

mask = res['frac_idiosyncratic'].values == '0.00%'
res.loc[mask, 'frac_idiosyncratic'] = ''

# Get nice names
res = res[['level', 'MSE', 'frac_uncond_bias', 'frac_aggregate', 'frac_idiosyncratic', 'frac_irreducible']]
res['level'] = res['level'].replace('consensus', 'Consensus')
res['level'] = res['level'].replace('individual', 'Individual')
res.rename(columns = {'frac_uncond_bias': 'Uncond.',
                      'frac_aggregate': 'Agg. Shocks',
                      'frac_idiosyncratic': 'Idio. Shocks',
                      'frac_irreducible': 'Irreducible Error',
                      'level': ''},
           inplace = True)

# Get LaTeX output with nice labels
latex_output = res.to_latex(index = False, column_format = 'lllllll')
latex_output = latex_output.replace('\\toprule', '\\toprule & & \multicolumn{3}{c}{Bias} & & \\\\ \\cmidrule(lr{0.5em}){3-5}')

# Save outpu
with open('{}/MSE_decomposition.tex'.format(output_folder), "w") as text_file:
    text_file.write(latex_output)
